#This module solves and returns the states and irrigation rates for a single evaluation of the mpc scheduler (open-loop simulation)
#Imports
from formulate_scheduler  import FormulateScheduler, SolveFormulation
import numpy as np

data_min = np.loadtxt("./weights_casestudy/data_min_scl_c.txt")
data_max = np.loadtxt("./weights_casestudy/data_max_scl_c.txt")

horizonLength = 20 #prediction horizone

#Obtain states from the optimization solution
def GetStateValues(optimalValues):
    predictedTrajectory =[]
    stateSeedPoint = 15
    
    for j in range(horizonLength):
        predictedTrajectory.append(optimalValues[int(stateSeedPoint+j*(5))])
        pass

    predictedTrajectory_unscaled = (np.array(predictedTrajectory) * (data_max[0] - data_min[0])) + data_min[0]
    return predictedTrajectory_unscaled

#Obtain states(unscaled) from the optimization solution
def GetStateValues_unscaled(optimalValues):
    
    predictedTrajectory =[]
    stateSeedPoint = 15
    
    for j in range(horizonLength):
        predictedTrajectory.append(optimalValues[int(stateSeedPoint+j*(5))])
        pass
    predictedTrajectory = np.array(predictedTrajectory)
    return predictedTrajectory

#Obtain irrigation rates from the optimization solution
def GetIrrigationAmount(optimalValues):
    prescribedIrrigation = []
    inputSeedPoint = 13
    
    for j in range(horizonLength):
        prescribedIrrigation.append(optimalValues[int(inputSeedPoint+j*(5))])
        pass
    return prescribedIrrigation

def GetIrrigationDecisions(optimalValues):
    irrigationDecisions = []
    decisionSeedPoint = 14
    
    for j in range(horizonLength):
        irrigationDecisions.append(optimalValues[int(decisionSeedPoint+j*(5))])
        pass
    return irrigationDecisions

def SolveScheduler(currentStates, previousInputs, cropCoeff, refEvap, lbx, ubx, guessX, lbSlack, 
                   ubSlack, QLower, QUpper, lowerZone, upperZone, Rmatrix, guessU, ubU,lbU, guessY, ubY, lbY, slope, guessesX, guessesU, guessesY):
    
    print("Open-loop MPC --- Start -------")    
    w_cont, J_cont, lbw_cont, ubw_cont, Guess_cont, G_cont, lbg_cont, ubg_cont = FormulateScheduler(currentStates, previousInputs, cropCoeff, refEvap, lbx, ubx, guessX, lbSlack, 
                   ubSlack, QLower, QUpper, lowerZone, upperZone, Rmatrix, guessU, ubU, lbU, guessY, ubY, lbY, slope, guessesX, guessesU, guessesY)
    
   
    w_opt_continuous = SolveFormulation(Guess_cont, w_cont, lbw_cont, ubw_cont, G_cont, J_cont, lbg_cont, ubg_cont)
    
    states = GetStateValues(w_opt_continuous)
    irrigationAmounts = GetIrrigationAmount(w_opt_continuous)
    irrigationDecisions = GetIrrigationDecisions(w_opt_continuous)
    
    print("Open-loop MPC --- END -------")
    
    return irrigationAmounts, irrigationDecisions, states


